#!/bin/bash
set -e

export LOG_AFTER_STEP_NUM="100"
export LOG_FILE="info.csv"

export ENABLE_INIT_ON_CPU=1
export XLA_FLAGS="--xla_backend_extra_options='xla_dtu_mixed_precision'"
export TF_XLA_FLAGS="--tf_xla_auto_jit=-1 --tf_xla_min_cluster_size=4"
export ENFLAME_AUTO4C=true

# python enflame_model_test.py --is_training=True \
#     --device=dtu --model=resnet --dtype=bf16 \
#     --data_format=CHNW --dataset=imagenet --batch_size=64 \
#     --epoch=90 --enable_saver=True --enable_evaluate=True \
#     --optimizer=sgd

python enflame_model_test.py --is_training=True \
    --device=dtu --model=resnet --dtype=fp16 \
    --data_format=CHNW --dataset=imagenet --batch_size=64 \
    --epoch=90 --enable_saver=True --enable_evaluate=True \
    --optimizer=sgd

# python enflame_model_test.py --is_training=True \
#     --device=dtu --model=resnet --dtype=fp32 \
#     --data_format=CHNW --dataset=imagenet --batch_size=16 \
#     --epoch=90 --enable_saver=True --enable_evaluate=True \
#     --optimizer=sgd


export ENABLE_INIT_ON_CPU=1
export TF_XLA_FLAGS="--tf_xla_auto_jit=-1 --tf_xla_min_cluster_size=4"
export DTU_UMD_FLAGS='ib_pool_size=134217728'
export ENFLAME_AUTO4C=true

# python enflame_model_test.py --is_training=True \
#     --device=dtu --model=inception_v3 --dtype=bf16 \
#     --data_format=CHNW --dataset=imagenet --batch_size=32 \
#     --epoch=90 --enable_saver=True --enable_evaluate=True \
#     --optimizer=sgd

python enflame_model_test.py --is_training=True \
    --device=dtu --model=inception_v3 --dtype=fp16 \
    --data_format=CHNW --dataset=imagenet --batch_size=32 \
    --epoch=90 --enable_saver=True --enable_evaluate=True \
    --optimizer=sgd


# python enflame_model_test.py --is_training=True \
#     --device=dtu --model=inception_v3 --dtype=fp32 \
#     --data_format=CHNW --dataset=imagenet --batch_size=16 \
#     --epoch=90 --enable_saver=True --enable_evaluate=True \
#     --optimizer=sgd




